package dslab.nameserver;

import java.io.InputStream;
import java.io.PrintStream;
import java.rmi.*;
import java.rmi.registry.LocateRegistry;
import java.rmi.registry.Registry;
import java.rmi.server.UnicastRemoteObject;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import at.ac.tuwien.dsg.orvell.Shell;
import at.ac.tuwien.dsg.orvell.StopShellException;
import at.ac.tuwien.dsg.orvell.annotation.Command;
import dslab.ComponentFactory;
import dslab.util.Config;
import dslab.util.MisconfigurationException;

public class Nameserver implements INameserver {

    // configuration keys that have to be present
    private static final String[] requiredConfigKeys = new String[] {
        "root_id", "registry.host", "registry.port"
    };

    private final String componentId;
    private final Config config;
    private final InputStream inputStream;
    private final PrintStream printStream;

    private final boolean root;
    private final String root_id;
    // own domain
    private final String domain;
    // nameserver stores its records by the following association
    // zone -> RMI remote
    private final Map<String, INameserverRemote> nameservers;
    private final Map<String, String> addresses;
    private final Registry registry;

    private INameserverRemote nameserverRemote;

    /**
     * Creates a new server instance.
     *
     * @param componentId the id of the component that corresponds to the Config resource
     * @param config the component config
     * @param in the input stream to read console input from
     * @param out the output stream to write console output to
     */
    public Nameserver(String componentId, Config config, InputStream in,
                      PrintStream out) throws MisconfigurationException, RemoteException {
        // print configuration for the benefit of the user
        System.out.printf("@%s: configuration:%n", componentId);
        for (String key : config.listKeys())
            System.out.printf("@%s: { %s: %s }%n", componentId, key, config.getString(key));

        this.componentId = componentId;
        this.config = config;
        this.inputStream = in;
        this.printStream = out;

        // check if required keys are present
        for (String key : requiredConfigKeys)
            if (!config.containsKey(key))
                throw new MisconfigurationException();

        // root nameserver does not have a domain, for starters we assume if it lacks the key
        // domain, it's a root server
        if (!config.containsKey("domain")) {
            this.domain = "";
            this.root = true;
        } else {
            this.domain = config.getString("domain");
            this.root = false;
        }
        this.root_id = config.getString("root_id");
        this.nameservers = new ConcurrentHashMap<>();
        this.addresses = new ConcurrentHashMap<>();

        // fixes weird issue where the JVM can't set the hostname correctly
        System.setProperty("java.rmi.server.hostname", config.getString("registry.host"));

        // root registry special treatment
        if (root)
            this.registry = LocateRegistry.createRegistry(config.getInt("registry.port"));
        else
            this.registry = LocateRegistry.getRegistry(config.getString("registry.host"), config.getInt("registry.port"));
        System.out.printf("@%s: registry %s%n", componentId, root ? "created as root" : "bound");
    }

    @Override
    public void run() {
        INameserverRemote root_stub = null;
        // try setting up RMI things
        try {
            if (!root) {
                root_stub = (INameserverRemote) registry.lookup(root_id);
                System.out.printf("@%s: got root server stub from registry%n", componentId);
            }

            // initialize nameserver remote (and let it export objects at runtime)
            nameserverRemote = new NameserverRemote(
                    String.format("%s-registration", componentId),
                    root,
                    config.getString("root_id"),
                    nameservers,
                    addresses,
                    registry,
                    root_stub
            );
            // export nameserver remote object
            INameserverRemote stub = (INameserverRemote) UnicastRemoteObject.exportObject(nameserverRemote, 0);

            // root binds his server stub to the registry using root_id
            if (root) {
                registry.bind(root_id, stub);
                System.out.printf("@%s: bound server stub \"%s\" to registry%n", componentId, root_id);
            }

            // non-root nameserver registers himself
            if (!root) {
                if (root_stub == null)
                    throw new NotBoundException();
                System.out.printf("@%s: attempting to register nameserver%n", componentId);
                root_stub.registerNameserver(domain, nameserverRemote);
                System.out.printf("@%s: registered nameserver%n", componentId);
            }
        } catch (NotBoundException | ConnectException exception) {
            System.err.printf("@%s: error: couldn't bind to registry%n", componentId);
            exception.printStackTrace();
            // can't continue
            shutdown();
        } catch (RemoteException | AlreadyRegisteredException | InvalidDomainException |
                 AlreadyBoundException exception) {
            System.err.printf("@%s: there was an error exporting the nameserver remote %s%n", componentId, domain);
            exception.printStackTrace();
            // can't continue
            shutdown();
        }

        // finally, instantiate shell
        Shell shell = new Shell(inputStream, printStream);
        shell.register(this);
        shell.setPrompt(String.format("%s> ", componentId));
        shell.run();
    }

    @Override
    @Command
    public void nameservers() {
        for (String zone : nameservers.keySet())
            System.out.printf("%s: %s%n", zone, nameservers.get(zone).toString());
    }

    @Override
    @Command
    public void addresses() {
        for (String zone : addresses.keySet())
            System.out.printf("%s: %s%n", zone, addresses.get(zone));
    }

    @Override
    @Command
    public void shutdown() throws StopShellException {
        // attempt un-exporting the exposed nameserver remote functions
        if (nameserverRemote != null) {
            try {
                UnicastRemoteObject.unexportObject(nameserverRemote, false);
                // try un-binding server stub as root
                if (root) {
                    registry.unbind(root_id);
                    System.out.printf("@%s: unbound server stub from registry%n", componentId);
                }
            } catch (NotBoundException | RemoteException exception) {
                System.out.printf("@%s: there was an error un-exporting the nameserver remote %s%n", componentId, domain);
                exception.printStackTrace();
            }
        }

        System.out.printf("@%s: stopped%n", componentId);

        // close shell
        throw new StopShellException();
    }

    public static void main(String[] args) throws Exception {
        INameserver component = ComponentFactory.createNameserver(args[0], System.in, System.out);
        component.run();
    }
}
